import hashlib
import os
import shutil
from pathlib import Path

import cv2
import gradio as gr
import torch

# 设置基础变量
model_file_path = 'yolov7.pt'
print(f'模型文件:{model_file_path}')

file_output_path = './runs/detect/videos/output.mp4'
print(f'视频输出文件:{file_output_path}')

# 界面标题配置
APP_TITLE = "xxx视频检测系统"
TAB_VIDEO = "视频处理"
TAB_IMAGE = "图片画廊"
TAB_SINGLE_IMAGE = "图片处理"
IMAGE_LABEL = "图片画廊"
INPUT_VIDEO_LABEL = "上传视频"
SINGLE_IMAGE_INPUT_LABEL = "上传图片"
SINGLE_IMAGE_OUTPUT_LABEL = "图片处理结果"
OUTPUT_VIDEO_LABEL = "预览结果"
PROCESS_IMG_LABEL = "处理图片"
IMG_PROCESS_ING = "图片处理中,请耐心等候!"
IMG_PROCESS_COMPLETE = "图片处理完成"
VIDEO_PROCESS_ING = "视频处理中,请耐心等候!"
VIDEO_PROCESS_COMPLETE = "视频处理完成"

# 加载模型
model = torch.hub.load('./', 'custom', model_file_path, force_reload=True, source='local', trust_repo=True)
if torch.cuda.is_available():
    # 如果有GPU可用，使用GPU
    device = torch.device("cuda")
    print("Using GPU")
else:
    # 否则，使用CPU
    device = torch.device("cpu")
    print("Using CPU")
# 设置设备
model.to(device)

# 判断项目是否存在runs/detect/images以及runs/detect/videos文件夹是否存在,如果不存在则创建
for dir_path in ['runs/detect/images', 'runs/detect/image', 'runs/detect/videos']:
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


# 清空文件夹
def clear_directory(directory_path):
    # 检查目录是否存在
    if os.path.exists(directory_path) and os.path.isdir(directory_path):
        # 遍历目录中的所有文件和子目录
        for filename in os.listdir(directory_path):
            file_path = os.path.join(directory_path, filename)
            try:
                # 如果是文件，则删除
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                # 如果是目录，则递归删除
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f'无法删除 {file_path}。原因: {e}')
    else:
        print(f'目录 {directory_path} 不存在。')


# 清空指定文件夹
def clear_all():
    clear_directory('runs/detect/images')
    clear_directory("runs/detect/videos")


# 获取图片集合
def get_images():
    return list(Path('runs/detect/images').glob('*.jpg'))


# 计算图像的哈希值,用于判断是否是同一张图片
def image_hash(image):
    return hashlib.md5(image.tobytes()).hexdigest()


# 处理单张图片
def single_process_image(image_file):
    # 清空图片文件夹
    clear_directory("runs/detect/image")
    # 模型处理图片
    results = model(image_file)
    # 保存图片并返回文件流
    results.save("runs/detect/image")
    return Path("runs/detect/image/image0.jpg")


# 处理上传的图片并预览结果
def process_image(video_file):
    if video_file is None:
        return None
    gr.Info(IMG_PROCESS_ING)
    # 打开视频文件
    cap = cv2.VideoCapture(video_file)
    # 创建输出目录
    captures_dir = Path('runs/detect/images')
    frame_id = 0
    last_hash = None
    # 每隔10帧保存一次图片
    skip_frames = 10

    # 逐帧处理
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 每隔10帧处理一次图片
        if frame_id % skip_frames == 0:
            results = model(frame)
            for *xyxy, conf, cls in results.xyxy[0]:
                if conf > 0.8:  # 假设置信度高于0.8为有效检测
                    x1, y1, x2, y2 = map(int, [xyxy[0], xyxy[1], xyxy[2], xyxy[3]])
                    # 计算当前帧的哈希值
                    current_hash = image_hash(frame[y1:y2, x1:x2])
                    # 如果当前帧的哈希值与上一帧的哈希值不同，则保存当前帧
                    if current_hash != last_hash:
                        save_path = captures_dir / f'frame_{frame_id}.jpg'
                        # 保存图片
                        cv2.imwrite(str(save_path), frame)
                        # 更新上一帧的哈希值
                        last_hash = current_hash
        frame_id += 1
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    gr.Info(IMG_PROCESS_COMPLETE)
    return get_images()


# 处理上传的视频并预览结果
def process_video(video_file, progress=gr.Progress()):
    progress(0, desc="Starting...")

    # 清空所有文件
    clear_all()

    gr.Info(VIDEO_PROCESS_ING)
    # 打开视频文件
    cap = cv2.VideoCapture(video_file)
    # 创建输出视频文件
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    # 获取视频的帧率、宽度和高度
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    # 创建输出视频对象
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fourcc = cv2.VideoWriter_fourcc('H', '2', '6', '4')
    # 创建输出视频文件
    out = cv2.VideoWriter(file_output_path, fourcc, fps, (width, height))
    # 获取视频的总帧数
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 逐帧处理
    while cap.isOpened():
        # 读取一帧
        ret, frame = cap.read()
        if not ret:
            break

        # 获取当前帧的帧号
        frame_id = cap.get(cv2.CAP_PROP_POS_FRAMES)
        # 更新进度条,保留两位小数
        percentage_value = round((frame_id / total_frames), 2)
        print(f"{percentage_value}%")
        # 更新progress中的文字内容
        progress(percentage_value)

        # 转换为RGB格式
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 模型预测
        results = model(frame_rgb)
        # 渲染结果
        frame_detected = results.render()[0]
        frame_detected_bgr = cv2.cvtColor(frame_detected, cv2.COLOR_RGB2BGR)
        # 写入输出文件
        out.write(frame_detected_bgr)
    # 释放资源
    cap.release()
    out.release()
    gr.Info(VIDEO_PROCESS_COMPLETE)
    return file_output_path


# 构建界面
with gr.Blocks(css="css/app.css", title=APP_TITLE) as app:
    gr.Markdown(APP_TITLE, elem_id="app_title")
    # 视频处理
    with gr.Tab(TAB_VIDEO):
        with gr.Row():
            with gr.Column(scale=5):
                input_video = gr.Video(elem_id="input_video", label=INPUT_VIDEO_LABEL)
            with gr.Column(scale=5):
                output_video = gr.Video(elem_id="output_video", label=OUTPUT_VIDEO_LABEL, interactive=False)
                btn_process_img = gr.Button(PROCESS_IMG_LABEL, elem_id="btn_process_img")
            input_video.change(process_video, inputs=[input_video], outputs=[output_video])

    # 图片画廊
    with gr.Tab(TAB_IMAGE):
        gallery = gr.Gallery(elem_id="image_list", label=IMAGE_LABEL)
    # 单张图片处理
    with gr.Tab(TAB_SINGLE_IMAGE):
        with gr.Row():
            with gr.Column(scale=5):
                input_image = gr.Image(elem_id="input_image", label=SINGLE_IMAGE_INPUT_LABEL)
            with gr.Column(scale=5):
                output_image = gr.Image(elem_id="output_image", label=SINGLE_IMAGE_OUTPUT_LABEL, interactive=False)
            input_image.change(single_process_image, inputs=[input_image], outputs=[output_image])

    # 图片处理
    btn_process_img.click(process_image, inputs=[output_video], outputs=[gallery])

# 启动应用
app.launch(server_name="127.0.0.1", server_port=7860)
